package com.swak.loadbalance.impl;

import java.util.List;
import java.util.Map;
import java.util.TreeMap;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

import com.swak.codec.Digests;
import com.swak.exception.DataAccessException;
import com.swak.loadbalance.Referer;

/**
 * 一致性hash 的负载均衡, 创建一个Hash环
 * 
 * @author lifeng
 * @date 2021年3月3日 下午12:52:35
 */
public class ConsistentHashLoadBalance<T> extends RandomLoadBalance<T> {

	private final ConcurrentMap<String, ConsistentHashSelector<?>> selectors = new ConcurrentHashMap<String, ConsistentHashSelector<?>>();

	@Override
	public void onRefresh(List<T> referers) {
		super.onRefresh(referers);
		this.selectors.clear();
	}

	@Override
	public T select(String key) {
		List<Referer<T>> referers = this.getReferers();
		if (referers == null) {
			throw new DataAccessException(this.getClass().getSimpleName() + " No available referers:");
		}
		Referer<T> ref = null;
		if (referers.size() > 1) {
			ref = doSelect(key);
		} else if (referers.size() == 1) {
			ref = referers.get(0);
		}

		if (ref != null) {
			return ref.select();
		}
		throw new DataAccessException(this.getClass().getSimpleName() + " No available referers");
	}

	@SuppressWarnings("unchecked")
	protected Referer<T> doSelect(String key) {
		List<Referer<T>> invokers = this.getReferers();
		// using the hashcode of list to compute the hash only pay attention to the
		// elements in the list
		int invokersHashCode = invokers.hashCode();
		ConsistentHashSelector<T> selector = (ConsistentHashSelector<T>) selectors.get(key);
		if (selector == null || selector.identityHashCode != invokersHashCode) {
			selectors.put(key, new ConsistentHashSelector<T>(invokers, invokersHashCode));
			selector = (ConsistentHashSelector<T>) selectors.get(key);
		}
		return selector.select(key);
	}

	private static final class ConsistentHashSelector<T> {

		private final TreeMap<Long, Referer<T>> virtualInvokers;

		private final int replicaNumber;

		private final int identityHashCode;

		ConsistentHashSelector(List<Referer<T>> invokers, int identityHashCode) {
			this.virtualInvokers = new TreeMap<Long, Referer<T>>();
			this.identityHashCode = identityHashCode;
			this.replicaNumber = 160;
			for (Referer<T> invoker : invokers) {
				String address = invoker.getKey();
				for (int i = 0; i < replicaNumber / 4; i++) {
					byte[] digest = Digests.md5((address + i).getBytes());
					for (int h = 0; h < 4; h++) {
						long m = hash(digest, h);
						virtualInvokers.put(m, invoker);
					}
				}
			}
		}

		public Referer<T> select(String key) {
			byte[] digest = Digests.md5(key.getBytes());
			return selectForKey(hash(digest, 0));
		}

		private Referer<T> selectForKey(long hash) {
			Map.Entry<Long, Referer<T>> entry = virtualInvokers.ceilingEntry(hash);
			if (entry == null) {
				entry = virtualInvokers.firstEntry();
			}
			return entry.getValue();
		}

		private long hash(byte[] digest, int number) {
			return (((long) (digest[3 + number * 4] & 0xFF) << 24) | ((long) (digest[2 + number * 4] & 0xFF) << 16)
					| ((long) (digest[1 + number * 4] & 0xFF) << 8) | (digest[number * 4] & 0xFF)) & 0xFFFFFFFFL;
		}
	}
}
